#!/usr/bin/env python
import os
import sys
import numpy as np
import pickle
import logging
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入自定义服务
try:
    from app.services.ml_service import SpamClassifier
except ImportError:
    logger.error("无法导入 SpamClassifier，请确保代码正确")
    sys.exit(1)

def prepare_training_data(csv_path=None):
    """
    准备训练数据，如果没有提供CSV，生成模拟数据
    """
    if csv_path and os.path.exists(csv_path):
        logger.info(f"从 {csv_path} 加载训练数据")
        try:
            data = pd.read_csv(csv_path)
            texts = data['content'].values
            labels = data['is_spam'].values  # 0或1
            return texts, labels
        except Exception as e:
            logger.error(f"加载CSV失败: {str(e)}")
    
    # 生成模拟数据
    logger.info("生成模拟训练数据")
    
    # 模拟正常内容（非垃圾）
    normal_texts = [
        "这篇文章写得很好，分析得很透彻",
        "我非常喜欢这个论坛的讨论氛围",
        "今天天气真好，很适合出去走走",
        "这个问题我思考了很久，终于有了答案",
        "感谢楼主分享这么有用的信息",
        "这个话题很有意思，值得深入讨论",
        "我对这个观点持保留意见，但尊重不同声音",
        "最近读了一本很好的书，推荐给大家",
        "这个建议非常实用，我已经采纳了",
        "论坛里的朋友们都很友善，谢谢大家",
        "我刚刚解决了这个技术问题，记录一下",
        "这次活动很成功，期待下次再参加",
        "周末去了公园，拍了很多照片，很放松",
        "这个软件的新功能真的很好用",
        "我认为这个政策有利有弊，需要辩证看待"
    ]
    
    # 模拟垃圾内容
    spam_texts = [
        "想要赚钱吗？快来加入我们的项目 http://money.make.com",
        "惊爆价促销，所有商品5折起，点击链接查看",
        "免费领取优惠券，添加微信好友13812345678",
        "我们提供最便宜的服务，联系电话10086",
        "一天轻松赚1000，点击了解详情 www.easy.money.com",
        "祛痘产品大促销，联系微信ABC123",
        "免费分享赚钱秘籍，添加QQ群123456789",
        "最新款手机批发价，联系电话获取报价",
        "急用钱？快速贷款，无抵押，当天放款",
        "在家轻松工作，月入过万，详情咨询",
        "暴利项目推广，月入10万，添加微信了解",
        "减肥产品大促销，买二送一，点击链接",
        "私人订制服务，价格优惠，联系电话",
        "免费领取试用装，仅需支付邮费9.9元",
        "特价机票预订，比官网便宜50%，微信咨询"
    ]
    
    # 合并数据
    all_texts = normal_texts + spam_texts
    all_labels = [0] * len(normal_texts) + [1] * len(spam_texts)
    
    # 创建数据文件目录
    os.makedirs(os.path.join('data', 'training'), exist_ok=True)
    
    # 保存为CSV备用
    df = pd.DataFrame({
        'content': all_texts,
        'is_spam': all_labels
    })
    csv_output_path = os.path.join('data', 'training', 'spam_sample.csv')
    df.to_csv(csv_output_path, index=False, encoding='utf-8')
    logger.info(f"模拟数据已保存到 {csv_output_path}")
    
    return np.array(all_texts), np.array(all_labels)

def train_spam_model():
    """训练垃圾内容分类模型并保存"""
    logger.info("开始训练垃圾内容分类模型")
    
    # 准备数据
    texts, labels = prepare_training_data()
    
    # 拆分训练和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        texts, labels, test_size=0.2, random_state=42
    )
    
    logger.info(f"训练数据大小: {len(X_train)}, 测试数据大小: {len(X_test)}")
    
    # 创建分类器
    classifier = SpamClassifier()
    
    # 确保向量化器和模型已创建
    if not classifier.vectorizer or not classifier.model:
        logger.error("分类器初始化失败")
        return False
    
    try:
        # 训练向量化器和分类器
        logger.info("拟合向量化器...")
        X_train_features = classifier.vectorizer.fit_transform(X_train)
        
        logger.info("训练分类器...")
        classifier.model.fit(X_train_features, y_train)
        classifier.is_trained = True
        
        # 评估模型
        logger.info("评估模型...")
        X_test_features = classifier.vectorizer.transform(X_test)
        y_pred = classifier.model.predict(X_test_features)
        
        accuracy = accuracy_score(y_test, y_pred)
        logger.info(f"模型准确率: {accuracy:.4f}")
        logger.info("\n" + classification_report(y_test, y_pred))
        
        # 创建模型目录
        os.makedirs(os.path.join('data', 'models'), exist_ok=True)
        
        # 保存模型
        model_path = os.path.join('data', 'models', 'spam_model.pkl')
        with open(model_path, 'wb') as f:
            pickle.dump({
                'vectorizer': classifier.vectorizer,
                'classifier': classifier.model
            }, f)
        
        logger.info(f"模型已保存到 {model_path}")
        return True
    except Exception as e:
        logger.error(f"训练失败: {str(e)}")
        return False

if __name__ == "__main__":
    logger.info("开始模型训练流程")
    
    # 训练垃圾内容分类模型
    if train_spam_model():
        logger.info("垃圾内容分类模型训练成功")
    else:
        logger.error("垃圾内容分类模型训练失败")
    
    logger.info("训练流程完成") 